#include "StateEstimator.h"
#include <string>
#include <cmath>
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

StateEstimator::StateEstimator()
    : a_com_B_(Vector3d()),
      a_com_O_(Vector3d()),
      omega_OB_B_(Vector3d()),
      omega_OB_O_(Vector3d()),
      omega_OBX_B_(Matrix3d()),
      p_i_B_(Vector3d()),
      p_dot_i_B_(Vector3d()),
      R_B_O_(Matrix3d()),
      R_B_imu_(Matrix3d()),
      R_imu_imu0_(Matrix3d()),
      R_imu0_O_(Matrix3d()),
      R_t0_imu_imu0_(Matrix3d()),
      x_k_(Vector<double, 18>())
{
}

StateEstimator::~StateEstimator()
{
}

void StateEstimator::setIMUdata(const Matrix3d& R_imu_imu0)
{
    R_imu_imu0_ = R_imu_imu0;
}

Matrix3d StateEstimator::getR_B_O()
{
    R_B_imu_ = R("x", 180.0);

    double psi_init = -asin(R_t0_imu_imu0_(1, 0)); // asin返回弧度
    psi_init = psi_init * 180.0 / 3.1415326;
    R_imu0_O_ = (R("z", psi_init)) * (R("x", 180));

    // 计算旋转矩阵
    R_B_O_ = R_imu0_O_ * R_imu_imu0_ * R_B_imu_;
    return R_B_O_;
}

Vector3d StateEstimator::geta_com_B(const Vector3d& a_original)
{
    Matrix3d R_imu_B_;
    R_B_imu_ = R("x", -180.0);

    a_com_B_ = R_imu_B_ * a_original;
    return a_com_B_;
}

Vector3d StateEstimator::getomega_OB_B(const Vector3d& omega_original)
{
    Matrix3d R_imu_B_;
    R_imu_B_ = R("x", -180.0);

    omega_OB_B_ = R_imu_B_ * omega_original;
    return omega_OB_B_;
}

Vector3d StateEstimator::getp_i_B(int i, const Vector3d& q_j_i)
{
    int ksi = (i == 2 || i == 4) ? 1 : -1;
    int del = (i == 1 || i == 2) ? 1 : -1;
    //(2.48)
    p_i_B_(0) = -L_2 * sin(q_j_i(1)) - L_3 * sin(q_j_i(1) + q_j_i(2)) + del * H_X;
    p_i_B_(1) = ksi * L_1 * cos(q_j_i(0)) + L_3 * sin(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) + L_2 * cos(q_j_i(1)) * sin(q_j_i(0)) + ksi * H_Y;
    p_i_B_(2) = ksi * L_1 * sin(q_j_i(0)) - L_3 * cos(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) - L_2 * cos(q_j_i(0)) * cos(q_j_i(1));
    return p_i_B_;
}

Vector3d StateEstimator::getp_dot_i_B(int i, const Vector3d& q_j_i, const Vector3d& q_dot_j_i)
{
    int ksi = (i == 2 || i == 4) ? 1 : -1;
    int del = (i == 1 || i == 2) ? 1 : -1;
    //(2.49)
    Matrix3d tmpmat;
    tmpmat(0, 0) = 0;
    tmpmat(0, 1) = -L_2 * cos(q_j_i(1)) - L_3 * cos(q_j_i(1) + q_j_i(2));
    tmpmat(0, 2) = -L_3 * cos(q_j_i(1) + q_j_i(2));
    tmpmat(1, 0) = ksi * L_1 * sin(q_j_i(0)) + L_3 * cos(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) + L_2 * cos(q_j_i(1)) * cos(q_j_i(0));
    tmpmat(1, 1) = -L_3 * sin(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2)) - L_2 * sin(q_j_i(1)) * sin(q_j_i(0));
    tmpmat(1, 2) = -L_3 * sin(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2));
    tmpmat(2, 0) = ksi * L_1 * cos(q_j_i(0)) + L_3 * sin(q_j_i(0)) * cos(q_j_i(1) + q_j_i(2)) + L_2 * cos(q_j_i(1)) * sin(q_j_i(0));
    tmpmat(2, 1) = L_3 * cos(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2)) + L_2 * sin(q_j_i(1)) * cos(q_j_i(0));
    tmpmat(2, 2) = L_3 * cos(q_j_i(0)) * sin(q_j_i(1) + q_j_i(2));

    p_dot_i_B_ = tmpmat * q_dot_j_i;

    return p_dot_i_B_;
}

Vector3d StateEstimator::geta_com_O()
{
    a_com_O_ = R_B_O_ * a_com_B_;
    return a_com_O_;
}

Vector3d StateEstimator::getomega_OB_O()
{
    omega_OB_O_ = R_B_O_ * omega_OB_B_;
    return omega_OB_O_;
}

Matrix3d StateEstimator::getomega_OBX_B()
{
    omega_OBX_B_ << 0, -omega_OB_B_(2), omega_OB_B_(1),
        omega_OB_B_(2), 0, -omega_OB_B_(0),
        -omega_OB_B_(1), omega_OB_B_(0), 0;
    return omega_OBX_B_;
}

Vector<double, 18> StateEstimator::getx_k(const Vector3d& q_j_i, const Vector3d& q_dot_j_i)
{
    getR_B_O();

    Vector<double, 28> z_k;
    // z_k(0-11)
    for (int i = 0; i < 4; i++)
    {
        z_k.segment(i * 3, 3) = -R_B_O_ * getp_i_B(i, q_j_i);
    }
    // z_k(12-23)
    for (int i = 0; i < 4; i++)
    {
        z_k.segment(12 + i * 3, 3) = -R_B_O_ * (getomega_OBX_B() * getp_i_B(i, q_j_i) + getp_dot_i_B(i, q_j_i, q_dot_j_i));
    }
    // z_k(24-27)
    for (int i = 0; i < 4; i++)
    {
        z_k(i + 24) = 0;
    }

    Matrix<double, 28, 18> H;
    // H(0-11)
    H.block<3, 3>(0, 0).setIdentity(3, 3);
    H.block<3, 3>(3, 0).setIdentity(3, 3);
    H.block<3, 3>(6, 0).setIdentity(3, 3);
    H.block<3, 3>(9, 0).setIdentity(3, 3);
    H.block<3, 3>(0, 3).setZero();
    H.block<3, 3>(3, 3).setZero();
    H.block<3, 3>(6, 3).setZero();
    H.block<3, 3>(9, 3).setZero();

    H.block<12, 12>(0, 6).setIdentity(12, 12);
    H.block<12, 12>(0, 6) = -H.block<12, 12>(0, 6);
    // H(12-23)
    H.block<3, 3>(12, 0).setZero();
    H.block<3, 3>(15, 0).setZero();
    H.block<3, 3>(18, 0).setZero();
    H.block<3, 3>(21, 0).setZero();
    H.block<3, 3>(12, 3).setIdentity(3, 3);
    H.block<3, 3>(15, 3).setIdentity(3, 3);
    H.block<3, 3>(18, 3).setIdentity(3, 3);
    H.block<3, 3>(21, 3).setIdentity(3, 3);

    H.block<12, 12>(0, 6).setZero();
    // H(24-27)
    H.block<4, 18>(24, 0).setZero();
    H(24, 8) = 1;
    H(25, 11) = 1;
    H(26, 14) = 1;
    H(27, 17) = 1;

    // colPivHouseholderQr分解求解线性方程组
    x_k_ = H.colPivHouseholderQr().solve(z_k);

    return x_k_;
}